import os
from os import path
from pathlib import Path
from typing import Optional

import cv2
import numpy as np
import pandas as pd
from torch import FloatTensor, tensor
from torch.utils.data import Dataset

from backend.model_core.global_variables import GlobalVariables
from backend.model_core.logger import logger


class FundusPreprocessor:
    """眼底图像预处理类"""

    def __init__(self, data_dir:Optional[str]=None, excel_path:Optional[str]=None):
        self.data_dir = data_dir
        self.excel_path = excel_path
        self.image_size = (512, 512)  # 添加图像尺寸配置

    def _validate_paths(self):
        """验证所有必要路径是否存在"""
        if self.data_dir and not path.exists(self.data_dir):
            raise ValueError(f"数据目录不存在: {self.data_dir}")

        if self.excel_path and not path.exists(self.excel_path):
            raise ValueError(f"Excel文件不存在: {self.excel_path}")

    def load_data(self):
        """加载数据集信息并转换性别字段"""
        self._validate_paths()

        try:
            # 使用原始字符串处理Windows路径
            df = pd.read_excel(Path(self.excel_path))

            # 新增性别编码转换
            df["Patient Sex"] = df["Patient Sex"].map({"Female": 0, "Male": 1})
            # 新增年龄标准化处理：转换为数值并计算最大年龄
            df["Patient Age"] = pd.to_numeric(df["Patient Age"], errors='coerce')
            print(df.head(5))
            logger.info(f"成功加载Excel文件: {self.excel_path}")
            return df
        except Exception as e:
            logger.error(f"加载Excel文件失败: {str(e)}")
            raise

    def _check_image_extension(self, image_path):
        """检查图像文件扩展名是否有效"""
        valid_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".tiff"]
        ext = path.splitext(image_path)[1].lower()
        if ext not in valid_extensions:
            raise ValueError(f"不支持的图像格式: {ext} 在路径: {image_path}")

    def preprocess_image(self, image_path):
        """单张图像预处理: 直方图均衡化 + 自适应对比度增强"""
        # 使用绝对路径确保加载成功
        full_path = path.join(self.data_dir, image_path.replace("/", os.sep))
        # 读入图片
        img = cv2.imread(full_path)

        if img is None:
            raise ValueError(f"无法加载图像: {full_path}（文件不存在或格式不支持）")
        ext = path.splitext(image_path)[1].lower()
        # 转换为灰度图（单通道）
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        if GlobalVariables.sample:
            # 保存灰度图
            gray_path = path.join(
                GlobalVariables.sample_save_dir,
                f"{path.basename(image_path)}_1_gray{ext}",
            )
            cv2.imwrite(gray_path, gray)

        # 自适应直方图均衡化
        clahe = cv2.createCLAHE(
            clipLimit=2.0, # 对比度（亮部与暗部的差异程度）限制阈值。值越大，对比度增强越明显，但可能产生噪声；值越小，增强效果越弱
            tileGridSize=(8, 8) # 表示将图像在水平方向和垂直方向上分别划分为 8 个块，因此总共会形成 8×8=64 个小区域（tiles）
        )
        enhanced = clahe.apply(gray)
        if GlobalVariables.sample:
            # 保存增强后的图像
            enhanced_path = path.join(
                GlobalVariables.sample_save_dir,
                f"{path.basename(image_path)}_2_enhanced{ext}",
            )
            cv2.imwrite(enhanced_path, enhanced)

        # 图像归一化
        normalized = enhanced / 255.0
        if GlobalVariables.sample:
            # 由于归一化后图像数据类型变为float32，范围在0-1之间，需要转换回uint8范围才能保存
            normalized_uint8 = (normalized * 255).astype(np.uint8)
            normalized_path = path.join(
                GlobalVariables.sample_save_dir,
                f"{path.basename(image_path)}_3_normalized{ext}",
            )
            cv2.imwrite(normalized_path, normalized_uint8)

        # 调整尺寸
        resized = cv2.resize(normalized, self.image_size)
        if GlobalVariables.sample:
            # 保存调整尺寸后的图像
            resized_uint8 = (resized * 255).astype(np.uint8)
            resized_path = path.join(
                GlobalVariables.sample_save_dir,
                f"{path.basename(image_path)}_4_resized{ext}",
            )
            cv2.imwrite(resized_path, resized_uint8)

        # `resized[..., np.newaxis]` 是 NumPy 中的切片和轴插入语法。
        # `...`（Ellipsis）表示选择所有未明确指定的维度。
        # `np.newaxis` 是一个特殊对象，用于在指定位置插入一个新的轴（维度）。
        # 这里的作用是在 `resized` 数组的最后添加一个新的通道维度。
        return resized[..., np.newaxis]  # 添加通道维度
    
    def create_dual_input(self, left_path, right_path):
        """创建双目输入: 拼接左右眼图像"""
        if GlobalVariables.mode == "test":
            GlobalVariables.test_current_pos += 1
        else:
            GlobalVariables.train_current_pos += 1
        current_pos = (
            GlobalVariables.test_current_pos
            if GlobalVariables.mode == "test"
            else GlobalVariables.train_current_pos
        )
        data_size = (
            GlobalVariables.test_data_size
            if GlobalVariables.mode == "test"
            else GlobalVariables.train_data_size
        )
        if (
            current_pos % GlobalVariables.batch_size == 0
            or current_pos == data_size
        ):
            logger.info(
                f"[{GlobalVariables.mode}] At Epoch {GlobalVariables.epoch} "
                f"Processing image {current_pos}"
                f"/{data_size}"
            )

        left = self.preprocess_image(left_path) # (512,512,1)
        right = self.preprocess_image(right_path) # (512,512,1)

        # 新增双目图像拼接保存
        if GlobalVariables.sample:
            # 构造保存路径
            ext = path.splitext(left_path)[1].lower()
            dual_path = path.join(
                GlobalVariables.sample_save_dir,
                f"{path.basename(left_path).split('.')[0]}_dual{ext}",
            )
            # 创建保存目录（如果不存在）
            os.makedirs(path.dirname(dual_path), exist_ok=True)
            # 水平拼接双通道图像并转换为单通道显示
            dual_display = np.hstack([left, right])
            # 转换为可保存的uint8格式
            dual_uint8 = (dual_display * 255).astype(np.uint8)
            # 保存拼接结果
            cv2.imwrite(dual_path, dual_uint8)

        return np.concatenate(
            [left, right], 
            axis=-1 # 在最后一个维度进行拼接
        )  # 形状: (512, 512, 2)


class FundusDataset(Dataset):
    """自定义眼底数据集"""

    def __init__(self, df, preprocessor):
        self.df = df
        self.preprocessor = preprocessor

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        image = self.preprocessor.create_dual_input(
            row["Left-Fundus"], row["Right-Fundus"]
        )
        
        # 添加维度转置 NHWC -> NCHW
        image = np.transpose(image, (2, 0, 1))  # 转置为(C,H,W)

        # 修改年龄和性别特征处理 - 增加年龄标准化
        age = tensor([row["Patient Age"] / GlobalVariables.max_age]).float()
        sex = tensor([row["Patient Sex"]]).float()

        # 修改为多标签输出格式
        multi_hot = tensor([row[col] for col in GlobalVariables.label_columns]).float()
        return {
            "image": FloatTensor(image),
            "age": age,
            "sex": sex,
            "labels": multi_hot,
        }
